import os
import operator
from dotenv import load_dotenv
from typing import Annotated, Any, TypedDict

from langchain_core.messages import AnyMessage, HumanMessage, AIMessage, ToolMessage
from langchain_core.tools import tool
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_openai import ChatOpenAI

from langgraph.graph import StateGraph, END
from langgraph.prebuilt import ToolNode
from langgraph.checkpoint.memory import MemorySaver

# --- 1. 环境与工具准备 ---
load_dotenv()

# a. 将工具分为"风险"和"安全"两类
# 风险工具：需要调用外部API，我们希望对其进行审核
try:
    risky_search_tool = TavilySearchResults(max_results=2)
except ImportError:
    from langchain_community.tools import DuckDuckGoSearchRun
    risky_search_tool = DuckDuckGoSearchRun()

# 安全工具：本地执行，我们信任它，无需审核
@tool
def python_calculator(expression: str) -> str:
    """一个可以执行 Python 数学表达式的计算器。"""
    try:
        result = eval(expression, {"__builtins__": None}, {})
        return f"表达式 '{expression}' 的计算结果是: {result}"
    except Exception as e:
        return f"执行表达式 '{expression}' 时出错: {e}"

safe_tools = [python_calculator]
risky_tools = [risky_search_tool]

# --- 2. 状态定义：增加循环次数计数器 ---
class AgentState(TypedDict):
    messages: Annotated[list[AnyMessage], operator.add]
    # 新增一个字段来跟踪执行的轮次
    iteration_count: int

# --- 3. 节点定义 ---
# a. Agent 节点 (大脑) - 保持不变
llm = ChatOpenAI(model="qwen-plus-latest", base_url=os.getenv("OPENAI_BASE_URL"))
# LLM 需要知道所有可用的工具，以便制定计划
model_with_tools = llm.bind_tools(safe_tools + risky_tools)

def agent_node(state: AgentState):
    print("\n---AGENT: 思考中...---")
    response = model_with_tools.invoke(state["messages"])
    return {"messages": [response]}

# b. 两个独立的工具节点 (双手)
safe_tool_node = ToolNode(safe_tools)
risky_tool_node = ToolNode(risky_tools)

# c. 智能路由节点 (决策中心)
MAX_ITERATIONS = 5
def router_node(state: AgentState) -> str:
    print("---ROUTER: 决策中...---")
    
    # 首先，检查是否超出最大循环次数
    iteration_count = state.get("iteration_count", 0) + 1
    if iteration_count > MAX_ITERATIONS:
        print(f"已达到最大循环次数 {MAX_ITERATIONS}，流程终止。")
        return END

    last_message = state["messages"][-1]

    if not isinstance(last_message, AIMessage) or not last_message.tool_calls:
        print("ROUTER: 无工具调用，流程结束。")
        return END

    # 根据将要调用的工具名称进行路由
    # 注意：这里简化为只检查第一个工具调用
    tool_name = last_message.tool_calls[0]["name"]
    print(f"ROUTER: 计划调用工具 '{tool_name}'")
    
    if tool_name in [t.name for t in safe_tools]:
        print("ROUTER: 路由到 -> 安全工具节点 (无需审批)")
        return "safe_tools"
    elif tool_name in [t.name for t in risky_tools]:
        print("ROUTER: 路由到 -> 风险工具节点 (需要审批)")
        return "risky_tools"
    else:
        print(f"ROUTER: 未知工具 '{tool_name}'，流程结束。")
        return END

# --- 4. 构建图 ---
workflow = StateGraph(AgentState)

workflow.add_node("agent", agent_node)
workflow.add_node("safe_tools", safe_tool_node)
workflow.add_node("risky_tools", risky_tool_node)

workflow.set_entry_point("agent")

# 使用新的智能路由节点来决定流程走向
workflow.add_conditional_edges(
    "agent",
    router_node,
    {
        "safe_tools": "safe_tools",
        "risky_tools": "risky_tools",
        END: END
    }
)

# 无论哪个工具节点执行完毕，都返回给 agent 继续思考
workflow.add_edge("safe_tools", "agent")
workflow.add_edge("risky_tools", "agent")

# --- 5. 编译图，精确设置中断点 ---
checkpointer = MemorySaver()
# 只在"风险工具"节点前中断
app = workflow.compile(checkpointer=checkpointer, interrupt_before=["risky_tools"])

# --- 6. 运行循环审批逻辑 ---
if __name__ == "__main__":
    question = "搜索一下Tavily和LangChain的GitHub星标数，然后计算它们的差值。"
    
    config = {"configurable": {"thread_id": "user_session_3"}}
    # 初始化状态，包括 iteration_count
    inputs = {
        "messages": [HumanMessage(content=question)],
        "iteration_count": 0
    }
    
    app.invoke(inputs, config)

    while True:
        snapshot = app.get_state(config)
        
        if snapshot.next:
            print("\n---流程暂停，等待您的审批---")
            print(f"下一步计划执行的节点: {snapshot.next[0]}")
            
            if snapshot.next[0] == "risky_tools":
                last_message = snapshot.values["messages"][-1]
                print("风险工具调用详情:", last_message.tool_calls)
            
            user_approval = input("您是否批准执行下一步？ (yes/no/exit): ")

            if user_approval.lower() == "yes":
                print("\n---用户批准，继续执行---")
                app.invoke(None, config=config)
                continue
            else:
                print("\n---用户拒绝或退出，流程终止---")
                break
        else:
            print("\n---流程已完成！---")
            final_message = snapshot.values["messages"][-1]
            print("最终答案:")
            final_message.pretty_print()
            break 